from typing import Dict, Optional, Sequence, Set, Tuple, overload

from onnx import AttributeProto

class SchemaError(Exception): ...

class OpSchema:
    @property
    def file(self) -> str: ...
    @property
    def line(self) -> int: ...
    @property
    def support_level(self) -> SupportType: ...
    @property
    def doc(self) -> Optional[str]: ...
    @property
    def since_version(self) -> int: ...
    @property
    def deprecated(self) -> bool: ...
    @property
    def domain(self) -> str: ...
    @property
    def name(self) -> str: ...
    @property
    def min_input(self) -> int: ...
    @property
    def max_input(self) -> int: ...
    @property
    def min_output(self) -> int: ...
    @property
    def max_output(self) -> int: ...
    @property
    def attributes(self) -> Dict[str, Attribute]: ...
    @property
    def inputs(self) -> Sequence[FormalParameter]: ...
    @property
    def outputs(self) -> Sequence[FormalParameter]: ...
    @property
    def type_constraints(self) -> Sequence[TypeConstraintParam]: ...
    @property
    def has_type_and_shape_inference_function(self) -> bool: ...
    @property
    def has_data_propagation_function(self) -> bool: ...
    @staticmethod
    def is_infinite(v: int) -> bool: ...
    def consumed(self, schema: OpSchema, i: int) -> Tuple[UseType, int]: ...
    def _infer_node_outputs(
        self,
        node_proto: bytes,
        value_types: Dict[str, bytes],
        input_data: Dict[str, bytes],
        input_sparse_data: Dict[str, bytes],
    ) -> Dict[str, bytes]: ...

    class TypeConstraintParam:
        @property
        def type_param_str(self) -> str: ...
        @property
        def description(self) -> str: ...
        @property
        def allowed_type_strs(self) -> Sequence[str]: ...

    class FormalParameterOption:
        Single: "OpSchema.FormalParameterOption" = ...
        Optional: "OpSchema.FormalParameterOption" = ...
        Variadic: "OpSchema.FormalParameterOption" = ...

    class DifferentiationCategory:
        Unknown: "OpSchema.DifferentiationCategory" = ...
        Differentiable: "OpSchema.DifferentiationCategory" = ...
        NonDifferentiable: "OpSchema.DifferentiationCategory" = ...

    class FormalParameter:
        @property
        def name(self) -> str: ...
        @property
        def types(self) -> Set[str]: ...
        @property
        def typeStr(self) -> str: ...
        @property
        def description(self) -> str: ...
        @property
        def option(self) -> OpSchema.FormalParameterOption: ...
        @property
        def isHomogeneous(self) -> bool: ...
        @property
        def differentiationCategory(self) -> OpSchema.DifferentiationCategory: ...

    class AttrType:
        FLOAT: "OpSchema.AttrType" = ...
        INT: "OpSchema.AttrType" = ...
        STRING: "OpSchema.AttrType" = ...
        TENSOR: "OpSchema.AttrType" = ...
        GRAPH: "OpSchema.AttrType" = ...
        SPARSE_TENSOR: "OpSchema.AttrType" = ...
        TYPE_PROTO: "OpSchema.AttrType" = ...
        FLOATS: "OpSchema.AttrType" = ...
        INTS: "OpSchema.AttrType" = ...
        STRINGS: "OpSchema.AttrType" = ...
        TENSORS: "OpSchema.AttrType" = ...
        GRAPHS: "OpSchema.AttrType" = ...
        SPARSE_TENSORS: "OpSchema.AttrType" = ...
        TYPE_PROTOS: "OpSchema.AttrType" = ...

    class Attribute:
        @property
        def name(self) -> str: ...
        @property
        def description(self) -> str: ...
        @property
        def type(self) -> OpSchema.AttrType: ...
        @property
        def default_value(self) -> AttributeProto: ...
        @property
        def required(self) -> bool: ...

    class SupportType(int):
        COMMON: "OpSchema.SupportType" = ...
        EXPERIMENTAL: "OpSchema.SupportType" = ...

    class UseType:
        DEFAULT: "OpSchema.UseType" = ...
        CONSUME_ALLOWED: "OpSchema.UseType" = ...
        CONSUME_ENFORCED: "OpSchema.UseType" = ...

def has_schema(op_type: str) -> bool: ...
def schema_version_map() -> Dict[str, Tuple[int, int]]: ...
@overload
def get_schema(
    op_type: str, max_inclusive_version: int, domain: str = ""
) -> OpSchema: ...
@overload
def get_schema(op_type: str, domain: str = "") -> OpSchema: ...
def get_all_schemas() -> Sequence[OpSchema]: ...
def get_all_schemas_with_history() -> Sequence[OpSchema]: ...
